#include <numeric>
#include <iostream>

#include "blas.hpp"

Tensor::Tensor()
{

}

Tensor::Tensor(std::vector<float> data)
{
    this->data_ = data;

    u_int32_t size = std::accumulate(this->shape_.begin(), this->shape_.end(), 1, std::multiplies<int>());

    for(u_int32_t i=0; i<size;i++)
    {
        this->data_.push_back(0);
    }
}

Tensor::Tensor(std::vector<u_int32_t> shape)
{
    this->shape_ = shape;

    u_int32_t size = std::accumulate(this->shape_.begin(), this->shape_.end(), 1, std::multiplies<int>());

    for(u_int32_t i=0; i<size;i++)
    {
        this->data_.push_back(0);
    }
}

void Tensor::view(std::vector<u_int32_t> shape)
{
    u_int32_t size = std::accumulate(shape.begin(), shape.end(), 1, std::multiplies<int>());
    if (size==this->getSize())
    {
        this->shape_ = shape;
    }
    else{
        std::cout << "error: expect size= " << this->getSize() <<", but get " << size << std::endl;
    }
}

std::vector<u_int32_t> Tensor::getShape() const
{
    return this->shape_;
}

std::vector<float> Tensor::getData() const
{
    return this->data_;
}

std::vector<float> Tensor::setData(std::vector<float> data)
{
    this->data_ = data;
    if (data.size() != this->getSize())
    {
        std::vector<u_int32_t> shape = {u_int32_t(data.size())};
        this->shape_ = shape;
    }
    return this->data_;
}

u_int32_t Tensor::getSize() const
{
    u_int32_t size = std::accumulate(this->shape_.begin(), this->shape_.end(), 1, std::multiplies<int>());
    return size;
}


// 重写 << 运算符
std::ostream &operator<<(std::ostream &stream, const Tensor &obj)
{
    stream << "Shape: ";
    for (int i = 0; i < obj.getShape().size(); i++)
    {
        stream << obj.getShape().at(i);
        if (i<obj.getShape().size()-1){
            stream << " x ";
        }
    }
    stream << " = " << obj.getSize() << std::endl;
    uint32_t max_print_num = 9;
    if (obj.getSize() > max_print_num)
    {
        for (int i = 0; i < max_print_num; i++)
        {
            stream << obj.getData().at(i) << " ";
        }
        stream  << std::endl;
        // stream << obj.getData().at(0) << " " << obj.getData().at(1) << " " << obj.getData().at(2) 
        // << " ... " << std::endl;
    }
    else
    {
        for (int i = 0; i < obj.getSize(); i++)
        {
            stream << obj.getData().at(i) << " ";
        }
        stream  << std::endl;
    }

    return stream;
};

Tensor Tensor::dot(Tensor tensor)
{
    if (this->shape_.back()==tensor.shape_.at(0))
    {
        std::vector<u_int32_t> out_shape = this->shape_;

        out_shape.erase(out_shape.end() - 1);
        out_shape.insert(out_shape.end(), tensor.shape_.begin()+1, tensor.shape_.end());

        Tensor out = Tensor(out_shape);

        uint32_t out_count = 0;
        u_int32_t tensor_step = std::accumulate(tensor.shape_.begin()+1, tensor.shape_.end(), 1, std::multiplies<u_int32_t>());

        for (uint32_t i=0; i < this->data_.size(); i += this->shape_.back())
        {
            for (uint32_t j=0; j < tensor_step; j++)
            {
                for (uint32_t k=0; k < this->shape_.back(); ++k)
                {
                    out.data_.at(out_count) += this->data_.at(i+k) * tensor.data_.at(j+tensor_step*k);
                }
                ++out_count;
            }
        }
        return out;
    }
    else
    {
        std::cout << "expect shape start with " << this->shape_.back() 
        << ", but got " << tensor.shape_.at(0) << std::endl;
        return Tensor();
    }
}
